"""猫狗识别 - Gradio 网页应用 (重构版)

此文件是对原有脚本的重构，目标是保留原有功能不变：
- 加载已训练好的集成模型、标准化器、PCA、Faiss 索引与训练标签
- 提供图片预处理、预测和基于 Faiss 的相似图片检索
- 使用 Gradio 提供交互式网页界面

文件路径常量与默认行为保持不变，控制台输出与返回文本也尽量保持一致。
"""

from typing import Tuple
import os
import warnings

import numpy as np
from PIL import Image
import joblib
import faiss
import gradio as gr

warnings.filterwarnings("ignore")

# -------------------- 全局配置 --------------------
IMAGE_SIZE = (224, 224)
MODEL_PATH = "./saved_models/ensemble_model.joblib"
SCALER_PATH = "./saved_models/scaler.joblib"
PCA_PATH = "./saved_models/pca.joblib"
INDEX_PATH = "./saved_indexes/faiss_index.faiss"
LABELS_PATH = "./saved_indexes/train_labels.npy"

# -------------------- 全局占位 (在 load_models 后会被赋值) --------------------
ensemble_model = None
scaler = None
pca = None
faiss_index = None
train_labels = None


def _file_exists(path: str) -> bool:
    return os.path.exists(path)


def load_models() -> Tuple[bool, str]:
    """加载模型、标准化器、PCA、Faiss 索引与训练标签。

    返回 (success, message)。保留原脚本的控制台输出以便行为一致。
    """
    global ensemble_model, scaler, pca, faiss_index, train_labels

    print("=" * 50)
    print("正在加载模型和索引...")
    print("=" * 50)

    try:
        # 依次检查所需文件
        required = [
            (MODEL_PATH, "模型文件不存在"),
            (SCALER_PATH, "标准化器文件不存在"),
            (PCA_PATH, "PCA文件不存在"),
            (INDEX_PATH, "Faiss索引文件不存在"),
            (LABELS_PATH, "训练标签文件不存在"),
        ]

        for path, msg in required:
            if not _file_exists(path):
                return False, f"错误: {msg} - {path}\n请先运行 train_ensemble.py 训练模型!"

        # 加载模型/器
        print("加载集成学习模型...")
        ensemble_model = joblib.load(MODEL_PATH)
        print(f"✓ 模型加载成功: {MODEL_PATH}")

        print("加载标准化器...")
        scaler = joblib.load(SCALER_PATH)
        print(f"✓ 标准化器加载成功: {SCALER_PATH}")

        print("加载PCA模型...")
        pca = joblib.load(PCA_PATH)
        print(f"✓ PCA模型加载成功: {PCA_PATH}")

        print("加载Faiss索引...")
        faiss_index = faiss.read_index(INDEX_PATH)
        print(f"✓ Faiss索引加载成功: {INDEX_PATH}")
        print(f"  索引中的向量数量: {faiss_index.ntotal}")

        print("加载训练标签...")
        train_labels = np.load(LABELS_PATH)
        print(f"✓ 训练标签加载成功: {LABELS_PATH}")
        print(f"  标签数量: {len(train_labels)}")

        print("=" * 50)
        print("所有模型和索引加载完成!")
        print("=" * 50)
        return True, "模型加载成功!"

    except Exception as exc:
        err = f"加载模型时发生错误: {str(exc)}"
        print(err)
        return False, err


# 程序启动时尝试加载
load_success, load_message = load_models()
if not load_success:
    print()
    print("!" * 50)
    print("警告: 模型加载失败!")
    print(load_message)
    print("!" * 50)


def preprocess_image(image) -> np.ndarray:
    """把输入图片转为模型需要的 2D numpy 特征向量 (1, N)。

    输入可以是 PIL.Image 或 numpy.ndarray（如 Gradio 提供的）。输出为 shape=(1, H*W*3) 的浮点数组，像原脚本那样将像素归一化到 [0,1] 并展平。
    """
    try:
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image.astype("uint8"))

        if image.mode != "RGB":
            image = image.convert("RGB")

        image = image.resize(IMAGE_SIZE)
        arr = np.array(image).astype("float32") / 255.0
        flat = arr.flatten().reshape(1, -1)
        return flat
    except Exception as exc:
        raise Exception(f"图片预处理失败: {str(exc)}")


def predict_image(image) -> str:
    """对单张图片进行预测并返回 Markdown 文本（与原脚本的返回格式保持一致）。"""
    # 与原脚本保持相同的前置检查与返回信息
    if ensemble_model is None or scaler is None or pca is None:
        return "❌ 错误: 模型未加载!\n\n请先运行 train_ensemble.py 训练模型,然后重启此应用。"

    if image is None:
        return "⚠️ 请先上传一张图片!"

    try:
        print()
        print("-" * 50)
        print("开始预测...")
        print("-" * 50)

        # 1. 预处理
        print("步骤1: 预处理图片...")
        processed = preprocess_image(image)
        print(f"✓ 预处理完成,特征维度: {processed.shape}")

        # 2. 标准化
        print("步骤2: 特征标准化...")
        scaled = scaler.transform(processed)
        print("✓ 标准化完成")

        # 3. PCA 降维
        print("步骤3: PCA降维...")
        pca_feat = pca.transform(scaled)
        print(f"✓ 降维完成,特征维度: {pca_feat.shape[1]}")

        # 4. 集成模型预测
        print("步骤4: 集成模型预测...")
        prediction = ensemble_model.predict(pca_feat)[0]

        if hasattr(ensemble_model, "predict_proba"):
            probs = ensemble_model.predict_proba(pca_feat)[0]
            cat_prob = probs[0] * 100
            dog_prob = probs[1] * 100
            print(f"✓ 预测完成 - 猫: {cat_prob:.2f}%, 狗: {dog_prob:.2f}%")
        else:
            cat_prob = 100.0 if prediction == 0 else 0.0
            dog_prob = 100.0 if prediction == 1 else 0.0
            print(f"✓ 预测完成 - 类别: {prediction}")

        # 5. Faiss 相似度检索（使用标准化后的向量与原脚本保持一致）
        print("步骤5: Faiss相似度检索...")
        if faiss_index is not None and train_labels is not None:
            query = scaled.astype(np.float32)
            k = 5
            distances, indices = faiss_index.search(query, k)
            neighbor_labels = train_labels[indices[0]]
            cat_neighbors = int(np.sum(neighbor_labels == 0))
            dog_neighbors = int(np.sum(neighbor_labels == 1))
            print(f"✓ 检索完成 - 最近{k}个邻居中: {cat_neighbors}个猫, {dog_neighbors}个狗")
        else:
            cat_neighbors = 0
            dog_neighbors = 0

        print("-" * 50)
        print("预测完成!")
        print("-" * 50)

        if prediction == 0:
            emoji = "🐱"
            label_text = "猫"
            confidence = cat_prob
        else:
            emoji = "🐶"
            label_text = "狗"
            confidence = dog_prob

        result = f"""
# 🎯 预测结果

## {emoji} 这是一只 **{label_text}**!

---

### 📊 预测详情

**置信度分布:**
- 🐱 猫: {cat_prob:.2f}%
- 🐶 狗: {dog_prob:.2f}%

**预测置信度:** {confidence:.2f}%

---

### 🔍 相似度分析 (Faiss检索)

在训练集中找到的最相似的5张图片:
- 🐱 猫: {cat_neighbors} 张
- 🐶 狗: {dog_neighbors} 张

---

### ✅ 预测完成
模型: 集成学习 (随机森林 + 逻辑回归)
特征: PCA降维 (500维)
"""

        return result

    except Exception as exc:
        error_result = f"""
# ❌ 预测失败

发生错误: {str(exc)}

请检查:
1. 图片格式是否正确 (支持 JPG, PNG 等常见格式)
2. 图片是否损坏
3. 模型文件是否完整

如果问题持续,请重新训练模型。
"""
        print(f"\n错误: {str(exc)}")
        return error_result


def create_interface() -> gr.Blocks:
    """构建并返回 Gradio 界面 (Blocks 对象)。"""
    custom_css = """
    .gradio-container { font-family: 'Arial', sans-serif; }
    .output-markdown { font-size: 16px; }
    """

    with gr.Blocks(css=custom_css, title="猫狗识别系统") as demo:
        gr.Markdown(
            """
        # 🐱🐶 猫狗识别系统

        基于集成学习和Faiss索引的智能图像分类系统

        ---

        ### 📝 使用说明
        1. 点击下方区域上传一张猫或狗的图片
        2. 系统会自动进行预测
        3. 查看预测结果和置信度

        ### 🔧 技术栈
        - **机器学习**: 随机森林 + 逻辑回归 (集成学习)
        - **特征降维**: PCA (主成分分析)
        - **相似度检索**: Faiss (Facebook AI Similarity Search)
        - **Web框架**: Gradio
        """
        )

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 上传图片")
                image_input = gr.Image(label="上传猫或狗的图片", type="pil", height=400)
                predict_button = gr.Button("🚀 开始预测", variant="primary", size="lg")

            with gr.Column(scale=1):
                gr.Markdown("### 📊 预测结果")
                output_text = gr.Markdown(value="等待上传图片...", label="预测结果")

        gr.Markdown(
            """
        ---
        ### 💡 提示
        - 支持的图片格式: JPG, PNG, BMP, WEBP 等
        - 图片会自动调整为 224x224 像素
        - 建议上传清晰的猫或狗的正面照片以获得最佳效果
        """
        )

        # 绑定交互（保持与原脚本行为一致：点击或上传时可触发预测）
        predict_button.click(fn=predict_image, inputs=image_input, outputs=output_text)
        image_input.change(fn=predict_image, inputs=image_input, outputs=output_text)

    return demo


if __name__ == "__main__":
    print()
    print("=" * 50)
    print("启动Gradio网页应用...")
    print("=" * 50)

    if not load_success:
        print()
        print("!" * 50)
        print("警告: 模型未成功加载!")
        print("应用将启动,但预测功能可能无法使用。")
        print("请先运行 train_ensemble.py 训练模型。")
        print("!" * 50)

    demo = create_interface()
    demo.launch(server_name="127.0.0.1", server_port=7860, share=False, show_error=True)
